package org.hit.burkun.mapping;

import java.io.IOException;
import java.io.InputStream;
import java.io.UnsupportedEncodingException;
import java.net.HttpURLConnection;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLConnection;
import java.net.URLEncoder;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import org.apache.commons.lang3.StringUtils;
import org.hit.burkun.mapping.IOnlineMappingProcessor.DBKEY;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class UniprotOnlineMapping {
	private static final String UNIPROT_SERVER = "http://www.uniprot.org/mapping/?";
	/**
	 * the key is retrieved from http://www.uniprot.org/help/programmatic_access
	 * 
	 * @author burkun
	 * 
	 */

	private static Logger logger = LoggerFactory
			.getLogger(UniprotOnlineMapping.class);

	private static String assembleRequestStr(DBKEY fromKey, DBKEY toKey,
			Collection<String> data) {
		StringBuilder locationBuilder = new StringBuilder(UNIPROT_SERVER);
		locationBuilder.append("from=");
		locationBuilder.append(fromKey.toString());
		locationBuilder.append("&to=");
		locationBuilder.append(toKey.toString());
		locationBuilder.append("&format=tab");
		locationBuilder.append("&query=");
		String postData = "";
		try {
			postData = URLEncoder.encode(StringUtils.join(data, " "), "utf-8");
		} catch (UnsupportedEncodingException e) {
			logger.error("Url encode wrong!");
		}
		locationBuilder.append(postData);
		logger.debug("Assemble request string:"
				+ locationBuilder.substring(
						0,
						locationBuilder.length() > 100 ? 100 : locationBuilder
								.length() - 1));
		return locationBuilder.toString();
	}

	public static void runMapping(DBKEY fromKey, DBKEY toKey,
			Collection<String> data, IOnlineMappingProcessor processor,
			int retryTimes) {
		logger.info("[START] run mapping form uniportKB...");
		String location = assembleRequestStr(fromKey, toKey, data);
		boolean isDone = false;
		while (!isDone && retryTimes > 0) {
			logger.debug("left " + retryTimes + " times to retry!");
			try {
				URL url = new URL(location);
				HttpURLConnection conn = (HttpURLConnection) url
						.openConnection();
				HttpURLConnection.setFollowRedirects(true);
				conn.setDoInput(true);
				conn.connect();
				int status = conn.getResponseCode();
				while (true) {
					int wait = 0;
					String header = conn.getHeaderField("Retry-After");
					if (header != null)
						wait = Integer.valueOf(header);
					if (wait == 0)
						break;
					logger.info("Waiting (" + wait + ")...");
					conn.disconnect();
					Thread.sleep(wait * 1000);
					conn = (HttpURLConnection) new URL(location)
							.openConnection();
					conn.setDoInput(true);
					conn.connect();
					status = conn.getResponseCode();
				}
				if (status == HttpURLConnection.HTTP_OK) {
					logger.info("Got a OK reply");
					InputStream reader = conn.getInputStream();
					URLConnection.guessContentTypeFromStream(reader);
					StringBuilder builder = new StringBuilder();
					int a = 0;
					while ((a = reader.read()) != -1) {
						builder.append((char) a);
					}
					// append
					Map<String, Set<String>> myMap = format(builder.toString());
					processor.process(fromKey, toKey, myMap);
					isDone = true;
				} else {
					retryTimes -= 1;
					logger.error("Failed, got " + conn.getResponseMessage()
							+ " for " + location);
				}
				conn.disconnect();
			} catch (MalformedURLException e) {
				retryTimes -= 1;
				logger.error(e.getMessage());
			} catch (IOException e) {
				retryTimes -= 1;
				logger.error(e.getMessage());
			} catch (InterruptedException e) {
				retryTimes -= 1;
				logger.error(e.getMessage());
			}
		}
		logger.info("[END] run mapping form uniportKB...");
	}

	private static Map<String, Set<String>> format(String res) {
		int beginIdx = 0;
		int endIdx = res.indexOf("\n");
		Map<String, Set<String>> myMap = new HashMap<String, Set<String>>();
		while (endIdx != -1) {
			String tempStr = res.substring(beginIdx, endIdx);
			beginIdx = endIdx + 1;
			endIdx = res.indexOf("\n", beginIdx);
			String[] terms = tempStr.split("\t");
			if (terms.length == 2 && !terms[0].equals("From")
					&& !terms[1].equals("To")) {
				if (myMap.containsKey(terms[0].trim())) {
					myMap.get(terms[0]).add(terms[1].trim());
				} else {
					HashSet<String> set = new HashSet<>();
					set.add(terms[1].trim());
					myMap.put(terms[0].trim(), set);
				}
			}
		}
		// 把最后的一个添加上
		String lastStr = res.substring(beginIdx);
		String[] terms = lastStr.split("\t");
		if (terms.length == 2 && !terms[0].equals("From")
				&& !terms[1].equals("To")) {
			if (myMap.containsKey(terms[0].trim())) {
				myMap.get(terms[0]).add(terms[1].trim());
			} else {
				HashSet<String> set = new HashSet<>();
				set.add(terms[1].trim());
				myMap.put(terms[0].trim(), set);
			}
		}
		return myMap;
	}
}
